Did you check for compatibility with torch.compile/jax.jit?

I generated a notebook to answer these questions (attached as PDF, sorry for LLM-braining
system-compile-jit.pdf
)

It shows:

• numpy / torch / jax construction and getters work

• extracted arrays work with torch.compile (fullgraph=True) and jax.jit

• constructing System inside torch.compile(..., fullgraph=True) fails on ctypes POINTER.new and ctypes graph-breaks to eager if fullgraph!=True

• jax.jit rejects System; constructing inside jit fails (dlpack) since metatomic.System is a host-side ctypes wrapper around mta_system_t, not a traced model type.

I'd say it's fine for now if we can not construct a system inside a torch.compile/jax.jit bit of code, but it would be nice to be able to compile a function like squared_sum_positions(system): return torch.sum(system.positions)**2. Looks like this works for torch, how hard would it be to make it work for jax? We don't have to do it in this PR, just make sure we are not painting ourself in a corner!

I do not understand the new requirement of releasing all arrays views before adding pairs/custom data. Can you explain?

The two APIs do not share buffers.
add_pairs only does system_mut.pairs.insert(...) (system.rs 371).
They're coupled because mta_system_get_data clones the Arc into manager_ctx (c_api/system.rs 191–209). add_pairs / add_custom_data go through Arc::get_mut (281–285 / 408–412), so any live getter view makes strong_count > 1, so mutation fails

We can keep this, or allow adds with live views (interior mutability on pairs / custom_data only). Which do you want?

On the topic of mutability, jax arrays strictly immutable. In torch, arrays may still look writable depending on version; the C comment on get_data is still “do not modify” (c_api/system.rs 168–170).
I would go for immutable, even if pytorch doesn't strictly require it.

Ah, right there is the Arc::get_mut thing. Ok, makes sense then!

On the topic of mutability, jax arrays strictly immutable. [...]

This should already be handled by DLPack, which has a flag for immutable data. So unless we gain something by specifying it ourself, I'd leave read-only considerations to DLPack and not mention it ourself.

I rebased and squashed all commits into a single one, but there are still some unresolved review comments here.

to pull the rebased changes:

git checkout python-wrap-system
git fetch
git reset --hard origin/python-wrap-system # (assuming the remote it called origin)

Rebased onto current metatomic-core, including #346. I kept your ManuallyDrop comments and the C++ test that a failed add_pairs / add_custom_data does not free the caller's system.

The review notes from the last pass are in this branch: backend detection uses isinstance, arrays_backend is a property, the getters no longer claim to be read-only (left to DLPack), System is documented before PairListOptions, pytest.raises matches the full message, and the refcount test builds its own system instead of using the fixture. Torch tests cover CPU, CUDA when it is available, and switching a torch system back to numpy. core-tests installs numpy and torch. I left jax out of tox, as you suggested.

After that I registered System as a JAX pytree, so jax.jit / value_and_grad can use positions and cell. Types, pbc, the length unit, and the backend stay static. Pair lists and custom data are rejected. Inside the transform there is no C pointer; it is created only if something asks for it once the arrays are concrete.

You said this did not have to land here, as long as system.positions stays a normal array. I am happy to drop that commit and do it as a follow-up. If it stays, two things should change: the jit test currently matches JAX's traceback-filter banner, which we do not control, and the tests are importorskip("jax"), so CI will not run them.

Written with an LLM assistant.